/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package code.google.pigfly.schema;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import org.apache.hadoop.hive.metastore.api.FieldSchema;
import org.apache.hadoop.hive.metastore.api.Schema;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector.Category;
import org.apache.hadoop.hive.serde2.typeinfo.ListTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.MapTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.PrimitiveTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.StructTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoUtils;

import code.google.pigfly.schema.PigflyFieldSchema.Type;

public class HCatSchemaUtils {

  private static HCatSchemaUtils ref = new HCatSchemaUtils();

  public static CollectionBuilder getStructSchemaBuilder() {
    return ref.new CollectionBuilder();
  }

  public static CollectionBuilder getListSchemaBuilder() {
    return ref.new CollectionBuilder();
  }

  public static MapBuilder getMapSchemaBuilder() {
    return ref.new MapBuilder();
  }

  public abstract class HCatSchemaBuilder {
    public abstract PigflySchema build() throws IOException;
  }

  public class CollectionBuilder extends HCatSchemaBuilder { // for
                                                             // STRUCTS(multiple-add-calls)
                                                             // and
                                                             // LISTS(single-add-call)
    List<PigflyFieldSchema> fieldSchemas = null;

    CollectionBuilder() {
      fieldSchemas = new ArrayList<PigflyFieldSchema>();
    }

    public CollectionBuilder addField(FieldSchema fieldSchema)
        throws IOException {
      return this.addField(getHCatFieldSchema(fieldSchema));
    }

    public CollectionBuilder addField(PigflyFieldSchema fieldColumnSchema) {
      fieldSchemas.add(fieldColumnSchema);
      return this;
    }

    @Override
    public PigflySchema build() throws IOException {
      return new PigflySchema(fieldSchemas);
    }

  }

  public class MapBuilder extends HCatSchemaBuilder {

    Type keyType = null;
    PigflySchema valueSchema = null;

    @Override
    public PigflySchema build() throws IOException {
      List<PigflyFieldSchema> fslist = new ArrayList<PigflyFieldSchema>();
      fslist.add(new PigflyFieldSchema(null, Type.MAP, keyType, valueSchema,
          null));
      return new PigflySchema(fslist);
    }

    public MapBuilder withValueSchema(PigflySchema valueSchema) {
      this.valueSchema = valueSchema;
      return this;
    }

    public MapBuilder withKeyType(Type keyType) {
      this.keyType = keyType;
      return this;
    }

  }

  /**
   * Convert a HCatFieldSchema to a FieldSchema
   * 
   * @param fs FieldSchema to convert
   * @return HCatFieldSchema representation of FieldSchema
   * @throws IOException
   */
  public static PigflyFieldSchema getHCatFieldSchema(FieldSchema fs)
      throws IOException {
    String fieldName = fs.getName();
    TypeInfo baseTypeInfo =
        TypeInfoUtils.getTypeInfoFromTypeString(fs.getType());
    return getHCatFieldSchema(fieldName, baseTypeInfo);
  }

  private static PigflyFieldSchema getHCatFieldSchema(String fieldName,
      TypeInfo fieldTypeInfo) throws IOException {
    Category typeCategory = fieldTypeInfo.getCategory();
    if (Category.PRIMITIVE == typeCategory) {
      return new PigflyFieldSchema(fieldName, getPrimitiveHType(fieldTypeInfo),
          null);
    } else if (Category.STRUCT == typeCategory) {
      PigflySchema subSchema =
          constructHCatSchema((StructTypeInfo) fieldTypeInfo);
      return new PigflyFieldSchema(fieldName, PigflyFieldSchema.Type.STRUCT,
          subSchema, null);
    } else if (Category.LIST == typeCategory) {
      PigflySchema subSchema =
          getHCatSchema(((ListTypeInfo) fieldTypeInfo).getListElementTypeInfo());
      return new PigflyFieldSchema(fieldName, PigflyFieldSchema.Type.ARRAY,
          subSchema, null);
    } else if (Category.MAP == typeCategory) {
      PigflyFieldSchema.Type mapKeyType =
          getPrimitiveHType(((MapTypeInfo) fieldTypeInfo).getMapKeyTypeInfo());
      PigflySchema subSchema =
          getHCatSchema(((MapTypeInfo) fieldTypeInfo).getMapValueTypeInfo());
      return new PigflyFieldSchema(fieldName, PigflyFieldSchema.Type.MAP,
          mapKeyType, subSchema, null);
    } else {
      throw new TypeNotPresentException(fieldTypeInfo.getTypeName(), null);
    }
  }

  private static Type getPrimitiveHType(TypeInfo basePrimitiveTypeInfo) {
    switch (((PrimitiveTypeInfo) basePrimitiveTypeInfo).getPrimitiveCategory()) {
    case BOOLEAN:
      return Type.BOOLEAN;
    case BYTE:
      return Type.TINYINT;
    case DOUBLE:
      return Type.DOUBLE;
    case FLOAT:
      return Type.FLOAT;
    case INT:
      return Type.INT;
    case LONG:
      return Type.BIGINT;
    case SHORT:
      return Type.SMALLINT;
    case STRING:
      return Type.STRING;
    default:
      throw new TypeNotPresentException(
          ((PrimitiveTypeInfo) basePrimitiveTypeInfo).getTypeName(), null);
    }
  }

  public static PigflySchema getHCatSchema(Schema schema) throws IOException {
    return getHCatSchema(schema.getFieldSchemas());
  }

  public static PigflySchema getHCatSchema(List<? extends FieldSchema> fslist)
      throws IOException {
    CollectionBuilder builder = getStructSchemaBuilder();
    for (FieldSchema fieldSchema : fslist) {
      builder.addField(fieldSchema);
    }
    return builder.build();
  }

  private static PigflySchema constructHCatSchema(StructTypeInfo stypeInfo)
      throws IOException {
    CollectionBuilder builder = getStructSchemaBuilder();
    for (String fieldName : ((StructTypeInfo) stypeInfo)
        .getAllStructFieldNames()) {
      builder.addField(getHCatFieldSchema(fieldName,
          ((StructTypeInfo) stypeInfo).getStructFieldTypeInfo(fieldName)));
    }
    return builder.build();
  }

  public static PigflySchema getHCatSchema(TypeInfo typeInfo)
      throws IOException {
    Category typeCategory = typeInfo.getCategory();
    if (Category.PRIMITIVE == typeCategory) {
      return getStructSchemaBuilder().addField(
          new PigflyFieldSchema(null, getPrimitiveHType(typeInfo), null))
          .build();
    } else if (Category.STRUCT == typeCategory) {
      PigflySchema subSchema = constructHCatSchema((StructTypeInfo) typeInfo);
      return getStructSchemaBuilder().addField(
          new PigflyFieldSchema(null, Type.STRUCT, subSchema, null)).build();
    } else if (Category.LIST == typeCategory) {
      CollectionBuilder builder = getStructSchemaBuilder();
      builder.addField(getHCatFieldSchema(null,
          ((ListTypeInfo) typeInfo).getListElementTypeInfo()));
      return builder.build();
    } else if (Category.MAP == typeCategory) {
      PigflyFieldSchema.Type mapKeyType =
          getPrimitiveHType(((MapTypeInfo) typeInfo).getMapKeyTypeInfo());
      PigflySchema subSchema =
          getHCatSchema(((MapTypeInfo) typeInfo).getMapValueTypeInfo());
      MapBuilder builder = getMapSchemaBuilder();
      return builder.withKeyType(mapKeyType).withValueSchema(subSchema).build();
    } else {
      throw new TypeNotPresentException(typeInfo.getTypeName(), null);
    }
  }

  public static PigflySchema getHCatSchemaFromTypeString(String typeString)
      throws IOException {
    return getHCatSchema(TypeInfoUtils.getTypeInfoFromTypeString(typeString));
  }

  public static PigflySchema getHCatSchema(String schemaString)
      throws IOException {
    if ((schemaString == null) || (schemaString.trim().isEmpty())) {
      return new PigflySchema(new ArrayList<PigflyFieldSchema>()); // empty
                                                                   // HSchema
                                                                   // construct
    }
    PigflySchema outerSchema =
        getHCatSchemaFromTypeString("struct<" + schemaString + ">");
    return outerSchema.get(0).getStructSubSchema();
  }

  public static FieldSchema getFieldSchema(PigflyFieldSchema hcatFieldSchema) {
    return new FieldSchema(hcatFieldSchema.getName(),
        hcatFieldSchema.getTypeString(), hcatFieldSchema.getComment());
  }

  public static List<FieldSchema> getFieldSchemas(
      List<PigflyFieldSchema> hcatFieldSchemas) {
    List<FieldSchema> lfs = new ArrayList<FieldSchema>();
    for (PigflyFieldSchema hfs : hcatFieldSchemas) {
      lfs.add(getFieldSchema(hfs));
    }
    return lfs;
  }
}
